import argparse
import os
import time

from app_config import AppConfig
from llm import call_llm
from llm import gen_subtitle_check_prompt
from log_manager import LogManager
from srt_formatter import format_srt


def check_subtitles(input, output_path, ref_path, model_name="llama2"):
    """
    生成字幕
    Args:
        video_path: 视频文件路径
        output_path: 输出的 SRT 文件路径
        model_name: Ollama 模型名称
    """
    # 读取生成的字幕文件
    with open(input, "r", encoding="utf-8") as f:
        subtitle_text = f.read()
    if ref_path:
        with open(ref_path, "r", encoding="utf-8") as f:
            ref_text = f.read()
    else:
        ref_text = None
    system_prompt, user_prompt = gen_subtitle_check_prompt(subtitle_text, ref_text)
    # 校对字幕
    corrected_subtitle = call_llm(system_prompt, user_prompt)
    LogManager.log("校对完成，开始格式化")
    # 格式化字幕
    improved_subtitles = format_srt(corrected_subtitle)
    LogManager.log("格式化完成，开始保存")
    # 保存优化后的字幕
    output_path = os.path.expanduser(output_path)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(improved_subtitles)
    LogManager.log(f"优化后的字幕已保存到:{output_path}")


if __name__ == "__main__":
    start_time = time.time()
    app_config = AppConfig()
    parser = argparse.ArgumentParser(description="视频字幕校验")
    parser.add_argument("input", help="字幕文件")
    parser.add_argument("--ref", default=None, help="原始文本文件路径")
    parser.add_argument("--output", default=None, help="输出的SRT文件路径")
    args = parser.parse_args()
    ref_path = args.ref
    input = args.input
    if args.output:
        output = args.output
    else:
        output = args.input.rsplit(".", 1)[0] + "_check.srt"
    check_subtitles(input, output, ref_path)
    end_time = time.time()
    total_seconds = end_time - start_time
    minutes = int(total_seconds // 60)
    seconds = total_seconds % 60
    LogManager.log(f"总运行时间: {minutes}分钟 {seconds:.2f}秒")
